Add UnboundedCountingSource::to for bounded reads from an UnboundedCountingSource - #39084
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the UnboundedCountingSource by adding the ability to specify an exclusive end limit for sequence generation. This change allows users to perform bounded reads from an unbounded source, which is a necessary step toward simplifying how GenerateSequence operates in streaming jobs. The implementation includes necessary updates to the source configuration, reader advancement logic, and watermark handling to ensure consistent behavior. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
…untingSource UnboundedCountingSource currently does not support configuring an end limit, this PR adds that capability. GenerateSequence with a from + to + rate boils down to a BoundedReadFromUnboundedSource. BoundedReadFromUnboundedSource splits emit all outputs from a single bundle that does not work well on Streaming jobs. UnboundedCountingSource::to is currently not exposed outside CountingSource. Eventually GenerateSequence can use UnboundedCountingSource::to instead of BoundedReadFromUnboundedSource
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an exclusive limit (end) to UnboundedCountingSource in CountingSource, allowing unbounded counting sources to stop at a specified value. It adds a .to(long end) builder method, updates serialization/equality methods, and modifies the reader to stop advancing and output a maximum watermark once the limit is reached. The review feedback correctly identifies a critical issue with the introduction of the mutable done state in UnboundedCountingReader. Since done is not checkpointed, restoring the reader from a checkpoint would cause watermark regressions. The reviewer suggests removing the done field entirely and instead deterministically calculating whether the reader is finished using current, source.stride, and source.end.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| private long current; | ||
| private boolean done = false; | ||
|
|
||
| // Initialized on first advance() |
There was a problem hiding this comment.
Instead of maintaining a mutable done state which is not checkpointed (and would cause watermark regressions upon reader restore), we can determine if the reader is done deterministically based on current and source.end. This allows us to completely remove the done field.
| private long current; | |
| private boolean done = false; | |
| // Initialized on first advance() | |
| private long current; | |
| // Initialized on first advance() |
| if (nextValue >= source.end) { | ||
| done = true; | ||
| return false; | ||
| } |
There was a problem hiding this comment.
Since we are removing the done field, we don't need to set done = true here. We can just return false directly when the next value reaches or exceeds the end limit.
if (nextValue >= source.end) {
return false;
}| @Override | ||
| public Instant getWatermark() { | ||
| return source.timestampFn.apply(current); | ||
| return done ? BoundedWindow.TIMESTAMP_MAX_VALUE : source.timestampFn.apply(current); |
There was a problem hiding this comment.
We can deterministically check if the reader is done by comparing current with source.end - source.stride. If current >= source.end - source.stride, then any subsequent value would be >= source.end, meaning no more elements can be produced. This avoids watermark regression when restoring from a checkpoint.
return (current >= source.end - source.stride)
? BoundedWindow.TIMESTAMP_MAX_VALUE
: source.timestampFn.apply(current);
a15ef42 to
a8bec67
Compare
…untingSource (apache#39084) * Add UnboundedCountingSource::to for bounded reads from an UnboundedCountingSource UnboundedCountingSource currently does not support configuring an end limit, this PR adds that capability. GenerateSequence with a from + to + rate boils down to a BoundedReadFromUnboundedSource. BoundedReadFromUnboundedSource splits emit all outputs from a single bundle that does not work well on Streaming jobs. UnboundedCountingSource::to is currently not exposed outside CountingSource. Eventually GenerateSequence can use UnboundedCountingSource::to instead of BoundedReadFromUnboundedSource
…untingSource (apache#39084) * Add UnboundedCountingSource::to for bounded reads from an UnboundedCountingSource UnboundedCountingSource currently does not support configuring an end limit, this PR adds that capability. GenerateSequence with a from + to + rate boils down to a BoundedReadFromUnboundedSource. BoundedReadFromUnboundedSource splits emit all outputs from a single bundle that does not work well on Streaming jobs. UnboundedCountingSource::to is currently not exposed outside CountingSource. Eventually GenerateSequence can use UnboundedCountingSource::to instead of BoundedReadFromUnboundedSource
2 participants
UnboundedCountingSource currently does not support configuring an end limit, this PR adds that capability.
GenerateSequence with a from + to + rate boils down to a BoundedReadFromUnboundedSource. BoundedReadFromUnboundedSource splits emit all outputs from a single bundle that does not work well on Streaming jobs.
UnboundedCountingSource::to is currently not exposed outside CountingSource. Eventually GenerateSequence can use UnboundedCountingSource::to instead of BoundedReadFromUnboundedSource